from VideoInfoDb import VideoInfoDb
import os
import websockets
import asyncio
import json
import threading
from manager import ManagerNode

import websockets.server
import time


class WebsocketsAPI(threading.Thread):
    def __init__(self, db_path):
        super().__init__()
        if not os.path.exists(db_path):
            os.makedirs(db_path)
        setting_json = os.path.join(db_path, "video_info.json")
        self.db = VideoInfoDb(setting_json)
        self.loop = asyncio.new_event_loop()
        self.stop_event = threading.Event()
        self.listen_server = None
        self.conditions = {}
        self.manager_ros = None
        self.generated_ply = False

    def set_manager_ros(self, manager_ros: ManagerNode):
        self.manager_ros = manager_ros

    async def websocket_callback(self, websocket, path):
        message = await websocket.recv()
        data = self.parse_message(message)

        api_type = data.get("type", "update_video_infos")
        handler = getattr(self, f"handle_{api_type}", self.handle_unknown)
        response = await handler(websocket, data)
        await websocket.send(json.dumps(response))

    def parse_message(self, message):
        try:
            return json.loads(message)
        except json.JSONDecodeError:
            return {}

    async def handle_default(self, websocket, data):
        return {"status": "error", "message": "No default action defined"}

    async def handle_update_video_infos(self, websocket, data):
        if not "video_path" in data or not "video_tags" in data:
            response = {"status": "error", "message": "Invalid data"}
            return response
        video_path = data["video_path"]
        video_tags = data["video_tags"]
        bag_path = data.get("bag_path", "")
        if bag_path is None:
            bag_path = ""
        try:
            video_time = video_path.split("/")[-1].split(".")[0]
            timestamp = time.strptime(video_time, "%Y-%m-%d-%H-%M-%S")
            timestamp = int(time.mktime(timestamp))
        except:
            print(f"Invalid video path: {video_path}")
            response = {"status": "error", "message": "Invalid video path"}
            return response
        self.db.update_tags(timestamp, video_path, video_tags, bag_path)
        response = {"status": "success", "message": "Received your data"}
        return response

    async def handle_get_ply(self, websocket, data):
        if not self.conditions.get("generated_ply", True):
            return {"status": "error", "message": "Already generated"}

        if self.manager_ros is None:
            return {"status": "error", "message": "ROS node not ready"}

        save_path = data.get("save_path")
        if not save_path:
            return {"status": "error", "message": "Invalid data"}

        self.conditions["generated_ply"] = False
        self.manager_ros.set_request("generate_ply", lambda: self.set_condition("generated_ply"), data)

        timeout = await self.wait_for_condition("generated_ply", 60)
        if not timeout:
            return {"status": "error", "message": "Timeout"}

        return {"status": "success", "message": "Received your data"}

    async def handle_reset(self, websocket, data):
        print("Resetting scene")
        if not self.conditions.get("reset", True):
            return {"status": "error", "message": "Already reset"}

        if self.manager_ros is None:
            return {"status": "error", "message": "ROS node not ready"}

        self.conditions["reset"] = False
        self.manager_ros.set_request("reset", lambda: self.set_condition("reset"), None)

        timeout = await self.wait_for_condition("reset", 30)
        if not timeout:
            return {"status": "error", "message": "Timeout"}

        return {"status": "success", "message": "Received your data"}

    async def handle_reset_random_seed(self, websocket, data):
        print("Resetting random seed")
        if not self.conditions.get("reset_random_seed", True):
            return {"status": "error", "message": "Already reset"}

        if self.manager_ros is None:
            return {"status": "error", "message": "ROS node not ready"}

        self.conditions["reset_random_seed"] = False
        self.manager_ros.set_request("reset_random_seed", lambda: self.set_condition("reset_random_seed"), data)

        timeout = await self.wait_for_condition("reset_random_seed", 30)
        if not timeout:
            return {"status": "error", "message": "Timeout"}

        return {"status": "success", "message": "Received your data"}

    async def handle_unknown(self, websocket, data):
        return {"status": "error", "message": "Unknown API type"}

    async def wait_for_condition(self, condition_name, timeout_seconds):
        timeout = 0
        while not self.conditions.get(condition_name, False):
            timeout += 1
            if timeout > timeout_seconds:
                self.conditions[condition_name] = True
                return False
            await asyncio.sleep(1)
        return True

    def set_condition(self, condition_name):
        self.conditions[condition_name] = True



    def run(self):
        asyncio.set_event_loop(self.loop)
        print("Start listening on ws://localhost:8765")
        self.listen_server = websockets.serve(
            self.websocket_callback, "localhost", 8765
        )
        self.loop.run_until_complete(self.listen_server)
        self.loop.run_forever()

    def stop(self):
        self.stop_event.set()
        self.db.close()
        if self.listen_server is not None:
            self.listen_server.ws_server.close()
        self.loop.call_soon_threadsafe(self.loop.stop)

    def close(self):
        self.db.close()


def main():
    api = WebsocketsAPI("video_info_db.json")
    api.start()
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        print("Stopping server...")
        api.stop()
        api.join()


if __name__ == "__main__":
    main()
